import time
import matplotlib.pyplot as plt
from torchsummary import summary
from torchvision.transforms._presets import ImageClassification
from PIL import Image
import os
from tqdm import tqdm
from torch.utils.data import DataLoader
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms
BATCH_SIZE = 16
# 设备配置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("设备：", device)

# 模型存储路径
# PATH_last = 'model/image_classification_last6_SGD10.pth'
# PATH_best = 'model/image_classification_best6_SGD10.pth'

transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor(),  # 转换为[C, H, W]
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 处理图片函数
def process_images(folder_path, label):
    results = []
    image_count = 0
    # 得到文件夹中的图片数量
    num_images = len([filename for filename in os.listdir(folder_path) if filename.lower().endswith(('.png', '.jpg', '.jpeg'))])
    # 遍历文件夹中的图片文件
    for filename in tqdm(os.listdir(folder_path)):
        if image_count >= num_images:
            break
        # 只处理图片文件
        if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
            try:
                # 加载并预处理图片
                image_path = os.path.join(folder_path, filename)
                image = Image.open(image_path).convert('RGB')
                tensor_image = transform(image)
                # 添加到结果列表，设置标签
                results.append((tensor_image, label))
                image_count += 1

            except Exception as e:
                print(f"处理图片 {filename} 时出错: {e}")

    return results

def test(valid_dataset, path):
    # 构建数据加载器
    dataloader = DataLoader(valid_dataset, batch_size=BATCH_SIZE, shuffle=True)
    # 加载模型并加载训练好的权重
    model = ImageClassification().to(device)  # 设备迁移
    model.load_state_dict(torch.load(path))
    model.eval()

    all_preds = []
    all_labels = []
    # 遍历每个batch的数据，获取预测结果，计算精度
    for x, y in dataloader:
        x, y = x.to(device), y.to(device)
        output = model(x)
        preds = torch.argmax(output, dim=-1)
        all_preds.extend(preds.cpu().numpy())
        all_labels.extend(y.cpu().numpy())

    # 计算正确率
    correct = sum(1 for p, l in zip(all_preds, all_labels) if p == l)
    accuracy = correct / len(all_preds)

    return accuracy  # 返回正确率


class ImageClassification(nn.Module):
    def __init__(self):
        super(ImageClassification, self).__init__()
        # 卷积层保持不变
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.pool1 = nn.MaxPool2d(2, 2)

        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.pool2 = nn.MaxPool2d(2, 2)

        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        self.pool3 = nn.MaxPool2d(2, 2)

        # conv1后：32x32 → 32x32x32 → 展平维
        # pool1后：32x32 → 16x16x32 → 展平维
        # conv2后：16x16 → 16x16x64 → 展平维
        # pool2后：16x16 → 8x8x64 → 展平维
        # conv3后：8x8 → 8x8x128 → 展平维
        # pool3后：8x8 → 4x4x128 → 展平维

        self.dropout = nn.Dropout(0.2)

        # 全连接层
        self.linear1 = nn.Linear(2048, 512)
        self.linear2 = nn.Linear(512, 128)
        self.linear3 = nn.Linear(128, 32)

        # 输出层
        self.out = nn.Linear(32, 10)

    def forward(self, x):
        # 卷积+relu+池化
        x = self.conv1(x)
        x = self.bn1(x)  # 先BN
        x = torch.relu(x)  # 再激活
        x = self.pool1(x)

        # 卷积+relu+池化
        x = self.conv2(x)
        x = self.bn2(x)  # BN
        x = torch.relu(x)
        x = self.pool2(x)

        # 卷积+relu+池化
        x = self.conv3(x)
        x = self.bn3(x)  # BN
        x = torch.relu(x)
        x = self.pool3(x)

        # 将特征图做成以为向量的形式：相当于特征向量
        x = x.reshape(x.size(0), -1)

        # 全连接层
        x = torch.relu(self.linear1(x))
        x = self.dropout(x)
        x = torch.relu(self.linear2(x))
        x = self.dropout(x)

        x = torch.relu(self.linear3(x))

        # 返回输出结果
        return self.out(x)


def train(model, train_dataset, valid_dataset, epoch, PATH_best, PATH_last, lr, opt='SGD'):
    acc=0
    criterion = nn.CrossEntropyLoss()

    print('优化器为--', opt)
    if opt == 'SGD':
        # SGD
        optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4)
    elif opt == 'Adam':
        # Adam
        optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=5e-4)
    else:
        print("请选择正确的优化器(SGD或Adam)")
        return

    # 学习率衰减
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.3)
    # 用于存储每个epoch的损失值
    losses = []
    # 用于存储每个epoch的验证集准确率
    accuracies = []
    for epoch_idx in range(epoch):
        dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
        sam_num = 0
        total_loss = 0.0
        start = time.time()
        for x, y in dataloader:
            x, y = x.to(device), y.to(device)
            output = model(x)
            loss = criterion(output, y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
            sam_num += 1

        avg_loss = total_loss / sam_num
        # 记录当前epoch的损失值
        losses.append(avg_loss)
        scheduler.step()
        # 保存模型
        torch.save(model.state_dict(), PATH_last)
        # 在验证集上测试模型
        accuracy = test(valid_dataset, path = PATH_last)
        if accuracy > acc:
            acc = accuracy
            torch.save(model.state_dict(), PATH_best)
        # 记录当前epoch的验证集准确率
        accuracies.append(accuracy)
        print('epoch:%2s loss:%.5f time:%.2fs accuracy:%2f' % (epoch_idx + 1, avg_loss, time.time() - start, accuracy))

    # 绘制损失曲线和准确率曲线
    plt.figure(figsize=(12, 5))
    plt.subplot(1, 2, 1)
    plt.plot(range(1, epoch + 1), losses, label='Training Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Training Loss Curve')
    plt.legend()

    plt.subplot(1, 2, 2)
    plt.plot(range(1, epoch + 1), accuracies, label='Validation Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.title('Validation Accuracy Curve')
    plt.legend()

    plt.show()


def pro_image(image_path):
    image = Image.open(image_path).convert('RGB')
    image = transform(image)
    return image

def predict_image(image_path, model_path):
    image = pro_image(image_path).unsqueeze(0).to(device)

    model = ImageClassification().to(device)
    model.load_state_dict(torch.load(model_path))
    model.eval()

    # 预测
    with torch.no_grad():
        output = model(image)
        probabilities = torch.softmax(output, dim=1)  # 转换为概率
        pred_class = torch.argmax(probabilities).item()  # 获取预测类别
        confidence = probabilities[0][pred_class].item()  # 获取置信度

    # 这里假设是0-9数字分类
    class_names = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']

    # 返回类别和置信度
    return class_names[pred_class], confidence


def train_model(epoch, batch_size, save_name, lr, opt):
    print('优化器为--', opt)
    PATH_best = 'model/' + save_name + '_beat' + '.pth'
    PATH_last = 'model/' + save_name + '_last' + '.pth'
    BATCH_SIZE = batch_size
    model = ImageClassification().to(device)
    summary(model, input_size=(3, 32, 32), batch_size=BATCH_SIZE, device=device.type)

    # 文件夹路径
    folder_path1 = "mnist_dataset/train/"
    folder_path2 = "mnist_dataset/test/"
    train_datasets = process_images(folder_path1 + '0', label=0)
    for i in range(1, 10):
        folder_path = folder_path1 + str(i)
        train_datasets += process_images(folder_path, label=i)

    valid_datasets = process_images(folder_path2 + '0', label=0)
    for i in range(1, 10):
        folder_path = folder_path2 + str(i)
        valid_datasets += process_images(folder_path, label=i)
    print('数据据加载完成')

    train(model, train_datasets, valid_datasets, epoch, PATH_best, PATH_last, lr = lr, opt=opt)

    print("训练完成")
    print("最优模型的正确率：")
    print(test(valid_datasets, path=PATH_best))



# if __name__ == '__main__':
#
#
#
#     train_model(10, batch_size = 16,  save_name = 'test1', lr = 0.01, opt = 'SGD')



    # print("预测图片：")




